from integrator import Integrator
from plugin_system import Plugin

import numpy as np

import copy

from simulation import NUMERIC_TYPE

class VerletIntegrator(Integrator):
    def __init__(self, timestep=None):
        Integrator.__init__(self, timestep=timestep)

    def integrate(self, mysolar, step):
        """ Verlet integrator
        Input:
          mysolar.positions : dictionary of { planet : positions  }
          mysolar.velocities : dictionary of { planet : velocities }
          mysolar.masses : dictionary of { planet : masses }
          dt : timestep in days
          nsteps : number of integration steps
          out : Output object with method callback()
        Output:
          new positions
          new velocities
        """
        positions_0 = mysolar.positions
        velocities_0 = mysolar.velocities
        masses = mysolar.masses
        kinetic_energy = self.kinetic_energy(velocities_0, mysolar.masses)

        # kick
        acceleration, pot_energy = self.accelerations(positions_0, mysolar.masses)
        velocities_1 = {}
        for planet in mysolar.planet_names:
            velocities_1[planet] = velocities_0[planet] + self.timestep * acceleration[planet]

        # drift
        positions_1 = {}
        for planet in mysolar.planet_names:
            positions_1[planet] = positions_0[planet] + self.timestep * velocities_1[planet]

        result = copy.deepcopy(mysolar)
        result.acceleration = acceleration
        result.positions = positions_1
        result.velocities = velocities_1
        result.kinetic_energy = kinetic_energy
        result.potential_energy = pot_energy
        result.simulated_time = step * self.timestep
        return result

class VerletIntegratorPluginAdapter(Plugin, VerletIntegrator):
    plugin_name = "VerletIntegrator"
    parameter_definitions = {"timestep" : (NUMERIC_TYPE, True, "Timestep")}
    
    def __init__(self, *args, **user_params):
        params = self.configure_parameters(**user_params)
        VerletIntegrator.__init__(self, timestep=params["timestep"])
        
